import React from 'react';
import { Tree as AntdTree } from 'antd';
import type { TreeProps as AntdTreeProps } from 'antd';
import { createStyles } from 'antd-style';

export interface TreeProps extends AntdTreeProps {
  /**
   * @description 自定义样式类名
   */
  customClassName?: string;
}

const useStyles = createStyles(({ token, css }) => ({
  tree: css`
    .ant-tree-node-selected {
      background-color: ${token.colorPrimaryBg};
    }
  `,
}));

const TreeComponent: React.FC<TreeProps> = ({ customClassName, className, ...restProps }) => {
  const { styles, cx } = useStyles();
  
  const treeClassName = cx(
    className,
    customClassName,
    styles.tree
  );

  return <AntdTree className={treeClassName} {...restProps as any} />;
};

const Tree = Object.assign(TreeComponent, {
  DirectoryTree: AntdTree.DirectoryTree,
  TreeNode: AntdTree.TreeNode,
});

export { Tree };
export default Tree;

